"""
增强版中文垃圾邮件检测器
包含更多功能和更好的性能
"""
import pandas as pd
import jieba.posseg as pseg
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
import re
import pickle

class ChineseSpamDetector:
    def __init__(self):
        self.vectorizer = None
        self.model = None
        self.stop_words = self._load_stop_words()
        
    def _load_stop_words(self):
        """加载中文停用词"""
        return set([
            '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到', '说', '要', '去',
            '你', '会', '着', '没有', '看', '好', '自己', '这', '那', '来', '用', '他', '她', '它', '们', '个', '啊', '吧', '呢',
            '吗', '嗯', '哦', '哈', '嘿', '呀', '哟', '呵', '唉', '呜', '哇', '咦', '为', '而', '如', '因', '所以', '但是', '然而',
            '可是', '虽然', '尽管', '无论', '不管', '无', '没', '还', '又', '才', '便', '这样', '那样', '这里', '那里', '这些', '那些'
        ])
    
    def preprocess_text(self, text):
        """高级中文文本预处理"""
        # 移除特殊字符，保留中文
        text = re.sub(r'[^\u4e00-\u9fff]', '', text)
        
        # 使用词性标注进行分词
        words = pseg.cut(text)
        
        # 只保留名词、动词、形容词等有意义的词
        meaningful_words = []
        for word, flag in words:
            if (len(word) >= 2 and 
                word not in self.stop_words and
                flag in ['n', 'v', 'a', 'nr', 'ns', 'nt', 'nz', 'vn', 'an']):
                meaningful_words.append(word)
        
        return ' '.join(meaningful_words)
    
    def create_extended_dataset(self):
        """创建扩展的中文数据集"""
        spam_emails = [
            "恭喜您！您已中奖100万元！立即拨打电话领取大奖！千万不要错过！",
            "免费获得减肥药！一周瘦20斤！效果保证！不瘦退款！立即订购！",
            "紧急通知！您的银行账户即将被冻结！请立即点击链接验证身份！",
            "投资理财！月收益30%！无风险！保本保息！立即加入！",
            "免费送苹果手机！限时优惠！点击领取！机会难得！",
            "代开发票！增值税普票专票！价格优惠！质量保证！",
            "网络兼职！在家赚钱！日入300元！无需经验！",
            "股票内幕消息！明日涨停板！千载难逢机会！立即购买！",
            "办理贷款！无抵押！当天放款！利息低！额度高！",
            "减肥神器！一个月瘦30斤！明星同款！效果惊人！",
            "免费试用化妆品！美白祛斑！立竿见影！包邮到家！",
            "彩票预测！百分百中奖！专业团队！成功率极高！",
            "赚钱机会！躺着就能赚钱！月入过万！立即加入！",
            "免费领取红包！无门槛！立即到账！点击领取！",
            "澳门赌场！网上百家乐！必赢秘籍！立即下载！",
            "办证刻章！身份证护照！质量保证！价格合理！",
            "股票推荐！内幕消息！稳赚不赔！立即购买！",
            "免费抽奖！豪车别墅！点击参与！中奖率极高！",
            "医疗广告！根治绝症！包治包好！立即咨询！",
            "投资项目！高回报！低风险！月收益50%！"
        ]
        
        normal_emails = [
            "明天下午3点开会，请大家准时参加。会议室在A栋302。",
            "您好，您的快递已经到达，请及时取件。快递单号：SF123456789。",
            "生日快乐！祝您身体健康，工作顺利，家庭幸福！",
            "会议纪要已发送，请查收附件并于明日前反馈意见。",
            "感谢您参加今天的培训，课件资料已上传到公司网盘。",
            "您的机票预订成功，航班号CA1234，明天早上8点起飞。",
            "项目进度汇报：目前已完成60%，预计下周完成全部工作。",
            "年终总结报告已提交，请领导审核批准。",
            "客户反馈很好，对我们的服务非常满意。",
            "明天的聚餐改到晚上7点，地点还是老地方。",
            "报销单据已整理完毕，请财务部门审核。",
            "产品发布会定于下周五举行，请提前准备相关材料。",
            "学术会议通知：计算机学会年会将于下月举行。",
            "图书馆新书到馆通知，欢迎大家前来借阅。",
            "公司年会节目征集，请大家积极报名参与。",
            "健康体检通知：请于本月内完成年度体检。",
            "工资单已发放，请查看邮箱附件。",
            "新员工入职培训安排在下周二上午。",
            "系统维护通知：明晚8点至12点系统将暂停服务。",
            "感谢您的订阅，我们会定期发送最新资讯。"
        ]
        
        emails = spam_emails + normal_emails
        labels = ['spam'] * len(spam_emails) + ['ham'] * len(normal_emails)
        
        return pd.DataFrame({'text': emails, 'label': labels})
    
    def train_with_multiple_models(self):
        """使用多种模型训练并选择最佳模型"""
        # 创建数据集
        data = self.create_extended_dataset()
        
        # 预处理
        data['processed_text'] = data['text'].apply(self.preprocess_text)
        
        # 特征提取
        self.vectorizer = TfidfVectorizer(
            max_features=2000,
            ngram_range=(1, 3),
            min_df=1,
            max_df=0.8
        )
        
        X = self.vectorizer.fit_transform(data['processed_text'])
        y = data['label'].map({'ham': 0, 'spam': 1})
        
        # 多模型比较
        models = {
            'Naive Bayes': MultinomialNB(),
            'SVM': SVC(kernel='linear', probability=True),
            'Random Forest': RandomForestClassifier(n_estimators=100)
        }
        
        best_score = 0
        best_model_name = ""
        
        print("模型性能比较：")
        for name, model in models.items():
            scores = cross_val_score(model, X, y, cv=5)
            avg_score = scores.mean()
            print(f"{name}: {avg_score:.4f} (+/- {scores.std() * 2:.4f})")
            
            if avg_score > best_score:
                best_score = avg_score
                best_model_name = name
                self.model = model
        
        print(f"\n最佳模型: {best_model_name}")
        
        # 训练最佳模型
        self.model.fit(X, y)
        
        return data
    
    def predict(self, text):
        """预测单个文本"""
        processed = self.preprocess_text(text)
        vector = self.vectorizer.transform([processed])
        prediction = self.model.predict(vector)[0]
        probability = self.model.predict_proba(vector)[0]
        
        return {
            'text': text,
            'processed': processed,
            'prediction': '垃圾邮件' if prediction == 1 else '正常邮件',
            'spam_probability': probability[1],
            'confidence': max(probability)
        }
    
    def save_model(self, filepath='chinese_spam_model.pkl'):
        """保存模型"""
        model_data = {
            'vectorizer': self.vectorizer,
            'model': self.model,
            'stop_words': self.stop_words
        }
        with open(filepath, 'wb') as f:
            pickle.dump(model_data, f)
        print(f"模型已保存到: {filepath}")
    
    def load_model(self, filepath='chinese_spam_model.pkl'):
        """加载模型"""
        with open(filepath, 'rb') as f:
            model_data = pickle.load(f)
        
        self.vectorizer = model_data['vectorizer']
        self.model = model_data['model']
        self.stop_words = model_data['stop_words']
        print("模型加载成功！")

def main():
    """主函数"""
    print("=== 中文垃圾邮件检测器 ===\n")
    
    # 创建检测器
    detector = ChineseSpamDetector()
    
    # 训练模型
    print("开始训练模型...\n")
    detector.train_with_multiple_models()
    
    # 保存模型
    detector.save_model()
    
    # 测试邮件
    test_emails = [
        "恭喜您中奖一千万！请立即领取！",
        "明天开会请准时参加。",
        "免费送手机！点击领取！",
        "您的快递已到达，请取件。",
        "投资理财！月息30%！无风险！",
        "会议纪要请查收附件。"
    ]
    
    print("\n=== 测试结果 ===")
    for email in test_emails:
        result = detector.predict(email)
        print(f"\n邮件: {result['text']}")
        print(f"预测: {result['prediction']}")
        print(f"垃圾邮件概率: {result['spam_probability']:.4f}")
        print(f"置信度: {result['confidence']:.4f}")

if __name__ == "__main__":
    main()